#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO人体检测模块
作者: 元歌
开发地: 广州南沙铁建中心
时间: 2024年10月
"""

import cv2
import numpy as np
from typing import List, Tuple, Optional, Dict
import os
from ultralytics import YOLO

class YOLODetector:
    """YOLO人体检测器 - 使用ultralytics库"""
    
    def __init__(self, 
                 model_path: str = None,
                 confidence_threshold: float = 0.5,
                 nms_threshold: float = 0.4,
                 target_class_id: int = 0):  # COCO数据集中person的class_id是0
        """
        初始化YOLO检测器
        
        Args:
            model_path: YOLO模型路径
            confidence_threshold: 置信度阈值
            nms_threshold: 非极大值抑制阈值
            target_class_id: 目标类别ID（person=0）
        """
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.target_class_id = target_class_id
        
        # 加载YOLO模型
        self._load_yolo_model(model_path)
    
    def _load_yolo_model(self, model_path: str = None):
        """加载YOLO模型"""
        try:
            if model_path is None:
                # 使用ultralytics默认的YOLOv8模型
                print("正在初始化YOLOv8n模型...")
                self.model = YOLO('yolov8n.pt')  # 自动下载nano版本
                print("YOLOv8n模型初始化完成")
            else:
                # 加载自定义模型
                if not os.path.exists(model_path):
                    raise FileNotFoundError(f"模型文件不存在: {model_path}")
                self.model = YOLO(model_path)
                print(f"成功加载自定义模型: {model_path}")
            
            # 设置模型参数
            self.model.overrides['conf'] = self.confidence_threshold
            self.model.overrides['iou'] = self.nms_threshold
            
        except Exception as e:
            print(f"YOLO模型加载失败: {e}")
            print("回退到HOG+SVM检测器")
            # 回退到HOG+SVM
            self.hog = cv2.HOGDescriptor()
            self.hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
            self.use_hog = True
            self.model = None
    
    def detect_humans(self, image: np.ndarray) -> List[Dict]:
        """
        检测图像中的人体
        
        Args:
            image: 输入图像
            
        Returns:
            检测到的人体列表，每个元素包含bbox和confidence
        """
        if self.model is not None:
            return self._detect_with_yolo(image)
        else:
            return self._detect_with_hog(image)
    
    def _detect_with_hog(self, image: np.ndarray) -> List[Dict]:
        """使用HOG+SVM检测人体"""
        try:
            # 多尺度检测
            (rects, weights) = self.hog.detectMultiScale(
                image,
                winStride=(8, 8),
                padding=(32, 32),
                scale=1.05,
                hitThreshold=0.0,
                useMeanshiftGrouping=False
            )
            
            detections = []
            for i, (x, y, w, h) in enumerate(rects):
                # 安全获取置信度
                if hasattr(weights, '__len__') and len(weights) > i:
                    confidence = weights[i][0] if hasattr(weights[i], '__len__') else weights[i]
                else:
                    confidence = 1.0
                
                if confidence > self.confidence_threshold:
                    detections.append({
                        'bbox': (x, y, w, h),
                        'confidence': confidence,
                        'class_id': 0,  # person
                        'class_name': 'person'
                    })
            
            return detections
        except Exception as e:
            print(f"HOG检测出错: {e}")
            return []
    
    def _detect_with_yolo(self, image: np.ndarray) -> List[Dict]:
        """使用YOLO检测人体"""
        try:
            # 使用ultralytics进行检测
            results = self.model(image, conf=self.confidence_threshold, iou=self.nms_threshold)
            
            detections = []
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        confidence = box.conf[0].cpu().numpy()
                        class_id = int(box.cls[0].cpu().numpy())
                        
                        # 只检测人体 (class_id = 0)
                        if class_id == self.target_class_id:
                            x, y, w, h = int(x1), int(y1), int(x2 - x1), int(y2 - y1)
                            detections.append({
                                'bbox': (x, y, w, h),
                                'confidence': float(confidence),
                                'class_id': class_id,
                                'class_name': 'person'
                            })
            
            return detections
        except Exception as e:
            print(f"YOLO检测出错: {e}")
            return []
    
    def draw_detections(self, image: np.ndarray, detections: List[Dict]) -> np.ndarray:
        """
        在图像上绘制检测结果
        
        Args:
            image: 输入图像
            detections: 检测结果列表
            
        Returns:
            绘制了检测框的图像
        """
        result_image = image.copy()
        
        for detection in detections:
            x, y, w, h = detection['bbox']
            confidence = detection['confidence']
            
            # 绘制边界框
            cv2.rectangle(result_image, (x, y), (x + w, y + h), (0, 255, 0), 2)
            
            # 绘制标签
            label = f"Person: {confidence:.2f}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
            cv2.rectangle(result_image, (x, y - label_size[1] - 10), 
                         (x + label_size[0], y), (0, 255, 0), -1)
            cv2.putText(result_image, label, (x, y - 5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
        
        return result_image
